package rsax

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
)

var (
	// ErrPrivateKeyDecode 私钥文件内容解码失败.
	ErrPrivateKeyDecode = errors.New("private key decode to block failed")

	// ErrPublicKeyDecode 公钥文件内容解码失败.
	ErrPublicKeyDecode = errors.New("public key decode to block failed")

	// ErrPublicKeyParse 公钥解析失败.
	ErrPublicKeyParse = errors.New("public key parsed failed")

	// ErrInvalidPublicKey 错误的公钥.
	ErrInvalidPublicKey = errors.New("invalid public key")
)

// PemPublicKey 二进制数据解析为公钥.
func PemPublicKey(key []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(key)
	if block == nil {
		return nil, ErrPublicKeyDecode
	}

	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, ErrPublicKeyParse
	}

	publicKey, ok := pubKey.(*rsa.PublicKey)
	if !ok {
		return nil, ErrInvalidPublicKey
	}

	return publicKey, nil
}

// PemPrivateKey 将二进制数据解析为私钥，支持pkcs1、pkcs8.
func PemPrivateKey(key []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(key)
	if block == nil {
		return nil, ErrPrivateKeyDecode
	}

	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		pkcs8PrivateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}

		privateKey, ok := pkcs8PrivateKey.(*rsa.PrivateKey)
		if !ok {
			return nil, ErrPrivateKeyDecode
		}

		return privateKey, nil
	}

	return privateKey, err
}
